#![feature(try_blocks)]

use anyhow::{Result, Context};
use clap::Clap;
use futures::StreamExt;
use futures::SinkExt;

use rand::Rng;
use byteorder::WriteBytesExt;

#[derive(Clap, Debug)]
#[clap(version = "1.0", author = "Donggu Kang")]
struct Opts {
    #[clap(long, default_value = "0.0.0.0:5555")]
    bind: String,
    // > leon: for instance, on 20200817 BTCUSD.20200925 product there are 5478639 book update
    #[clap(long, default_value = "15")] // our frequent feed has 15ms average interval
    ping_interval_ms: u64,
    #[clap(long, default_value = "5000")]
    msg_size: usize,
}

/*
Okex feed average msg size < 2K
To be safe, we use a payload 5K
*/

async fn accept_connection(stream: tokio::net::TcpStream, msg_size: usize, ping_interval_ms: u64) -> Result<()> {
    anyhow::ensure!(msg_size >= 20, "msg_size should be >= 20");

    stream.set_nodelay(true)?;

    let mut payload: Vec<u8> = rand::thread_rng()
        .sample_iter(&rand::distributions::Uniform::from(0..=255_u8))
        .take(msg_size)
        .collect();

    let addr = stream.peer_addr()?;
    let result: Result<()> = try {
        let ws_stream = tokio_tungstenite::accept_async(stream)
            .await
            .context("Error during the websocket handshake")?;

        let (mut write, _) = ws_stream.split();
        let mut interval = tokio::time::interval(std::time::Duration::from_millis(ping_interval_ms));
        loop {
            // send the current time
            let rt:i64 = chrono::Utc::now().timestamp_nanos();

            {
                let mut cursor = std::io::Cursor::new(&mut payload);
                cursor.write_i64::<byteorder::BigEndian>(rt)?;
            }

            // write.send(tungstenite::Message::Text(rt.to_string())).await.context("Fail to send ping msg")?;
            write.send(tungstenite::Message::Binary(payload.clone())).await.context("Fail to send ping msg")?;
            interval.tick().await;
        }
    };
    if let Err(e) = result {
        eprintln!("Error from the client {}", addr);
        eprintln!("{}", e);
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    let opts = Opts::parse();
    println!("{:?}", opts);

    let mut listener = tokio::net::TcpListener::bind(&opts.bind).await.context("Failed to start server. Check if your port is available.")?;

    while let Ok((stream, _)) = listener.accept().await {
        tokio::spawn(accept_connection(stream, opts.msg_size, opts.ping_interval_ms));
    }
    Ok(())
}
